package jtg.generator;

import com.microsoft.z3.ArithExpr;
import com.microsoft.z3.BoolExpr;
import com.microsoft.z3.Context;
import jtg.path.PrimePathCal;
import jtg.solver.Z3Solver;
import jtg.util.CommonUtil;
import jtg.util.PathUtil;
import soot.*;
import soot.jimple.ConditionExpr;
import soot.jimple.IfStmt;
import soot.jimple.IntConstant;
import soot.jimple.StaticFieldRef;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIdentityStmt;
import soot.jimple.internal.JIfStmt;
import soot.toolkits.graph.BriefUnitGraph;
import soot.toolkits.graph.UnitGraph;
import soot.util.Chain;

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;

public class MCDCCovGenerator extends GeneralGenerator {
    private static final Integer LOOP_MAX_CNT_IN_EXTEND = 2;
    private double coverageRate;
    private HashMap<Integer, List<Unit>> decisions = new HashMap<>();
    private String[] splitContent;

    public MCDCCovGenerator(String classPath, String className, String methodName, String javaPath) throws IOException {
        super(classPath, className, methodName);
        getDecisions();
        getSplitContent(javaPath);
    }
    //以$开头的变量是Jvars
    protected List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }
    //方法参数
    protected List<Local> getParameter() {
        return new ArrayList<Local>(body.getParameterLocals());
    }
    protected List<Local> getLocalVars() {
        List<Local> jVars = getJVars();
        List<Local> parameter = getParameter();
        Chain<Local> locals = body.getLocals();
        List<Local> localVars = new ArrayList<>();
        for (Local local : locals) {
            if (jVars.contains(local))
                continue;
            if (parameter.contains(local))
                continue;
            localVars.add(local);
        }
        return localVars;
    }

    public List<String> MCDCGenerate() throws Exception {
        ArrayList<String> sootTestSet = new ArrayList<String>();
        HashMap<Integer, List<Unit>> decisionsCopy = new HashMap<>(decisions);
        List<Local> jVars = getJVars();
        List<Local> localVars = getLocalVars();
        List<Local> paraVars = getParameter();
        localVars.addAll(jVars);
        int coveredCnt = 0;
        int conditionCnt = 0;
        for (int line : decisions.keySet()) {
            conditionCnt += (decisions.get(line).size());
        }

        for (Unit head : ug.getHeads()) {
            for (Unit tail : ug.getTails()) {
                if (decisionsCopy.size() == 0) break;
                //gat all complete paths
                ArrayList<List<Unit>> paths = new ArrayList<>();
                ArrayList<Unit> tmpPath = new ArrayList<>();
                tmpPath.add(head);
                getAllPathsHead2Tail(paths, tmpPath, new HashMap<>(), tail);
                for (List<Unit> path : paths) {
                    if (decisionsCopy.size() == 0) break;
                    HashMap<String, String> assignList = new HashMap<>();
                    ArrayList<String> stepConditions = new ArrayList<>();

                    //go through all units
                    for (Unit unit : path) {
                        if (decisionsCopy.size() == 0) break;

                        if (unit instanceof JAssignStmt) {
                            String leftOp = ((JAssignStmt) unit).getLeftOp().toString();
                            String rightOp = ((JAssignStmt) unit).getRightOp().toString();
                            Value rightOp1 = ((JAssignStmt) unit).getRightOp();

                            if (leftOp.contains(".")) {
                                leftOp = leftOp.substring(0, leftOp.indexOf(".")) + "_" + leftOp.substring(leftOp.lastIndexOf(" ") + 1, leftOp.length() - 1);
                            }
                            if (rightOp.contains(".")) {
                                rightOp = rightOp.substring(0, rightOp.indexOf(".")) + "_" + rightOp.substring(rightOp.lastIndexOf(" ") + 1, rightOp.length() - 1);
                            }
                            if (leftOp.contains("length")) {
                                leftOp = leftOp.substring(leftOp.indexOf(" ") + 1, leftOp.indexOf("("));
                            }
                            //字符串的length
                            if (rightOp.contains("length()")) {
                                rightOp = rightOp.substring(rightOp.indexOf(" ") + 1, rightOp.indexOf("("));
                            }
                            //i0 = lengthof r1数组
                            if (rightOp.contains("lengthof")) {
                                rightOp = rightOp.substring(rightOp.indexOf(" ") + 1) + "_" + "length";
                            }
                            //数组
                            if (rightOp.contains("[")) {
                                rightOp = rightOp.substring(0, 2) + "_" + rightOp.substring(rightOp.indexOf("[") + 1, (rightOp.indexOf("[") + 1 + 2));
                            }
                            //枚举
                            if (rightOp1 instanceof StaticFieldRef) {
                                rightOp1 = (StaticFieldRef) rightOp1;
                                String rightOpStr = rightOp1.toString();
                                Class<?> aClass = Class.forName(rightOp1.getType().toString());
                                if (aClass.isEnum()) {
                                    Class<Enum> enumClass = (Class<Enum>) aClass;
                                    Enum[] enumConstants = enumClass.getEnumConstants();
                                    for (Enum enumConstant : enumConstants) {
                                        String name = enumConstant.name();
                                        if (rightOpStr.contains(name)) {
                                            rightOpStr = String.valueOf(enumConstant.ordinal());
                                            break;
                                        }

                                    }
                                    rightOp = rightOpStr;
                                    //stepConditions.add();
                                }

                            }

                            for (Local paraVar : paraVars) {
                                if (rightOp.contains(paraVar.toString())) {
                                    for (String key : assignList.keySet()) {
                                        if (rightOp.contains(key)) {
                                            rightOp = rightOp.replace(key, assignList.get(key));
                                        }
                                    }
                                }
                            }
                            for (Local localVar : localVars) {
                                if (rightOp.contains(localVar.toString()) && assignList.get(localVar.toString()) != null) {
                                    rightOp = rightOp.replace(localVar.toString(), assignList.get(localVar.toString()));
                                }
                            }
                            assignList.put(leftOp, rightOp);
                        }
                        if (unit instanceof JIfStmt) {

                            int lineNumber = unit.getJavaSourceStartLineNumber();

                            if (decisionsCopy.containsKey(lineNumber)) {

                                ArrayList<String> currentStepConditions = new ArrayList<>(stepConditions);

                                ArrayList<String> conditions = new ArrayList<>();
                                List<Unit> sootConditions = decisionsCopy.get(lineNumber);
                                for (Unit sootCondition : sootConditions) {
                                    String condition = ((JIfStmt) sootCondition).getCondition().toString();
                                    //替换条件里的局部变量
                                    for (Local localVar : localVars) {
                                        if (condition.contains(localVar.toString()) && assignList.get(localVar.toString()) != null) {
                                            condition = condition.replace(localVar.toString(), assignList.get(localVar.toString()));
                                            //stopFlag= false;
                                        }
                                    }
                                    conditions.add(condition);
                                }
                                //deal the same decision once
                                decisionsCopy.remove(lineNumber);

                                System.out.println("lineNumber: " + lineNumber);
                                System.out.println("currentStepConditions: " + currentStepConditions);
                                System.out.println("conditions: " + conditions);

                                String logicalExpression = getLogicalExpressionByLine(lineNumber);
                                System.out.println("logicalExpression: " + logicalExpression);

                                String[] logicalExpressionSplit = logicalExpression.split("&&|\\|\\|");
                                for (int i = 0; i < logicalExpressionSplit.length; i++) {
                                    String s = logicalExpressionSplit[i].trim();
                                    if (s.charAt(0) == '(') {
                                        s = s.substring(1);
                                    }
                                    if (s.charAt(s.length() - 1) == ')') {
                                        s = s.substring(0, s.length() - 1);
                                    }
                                    logicalExpressionSplit[i] = s;
                                }

                                String[] operations = new String[logicalExpressionSplit.length];
                                int ptr = 0;
                                for (int i = 0; i < logicalExpression.length(); i++) {
                                    if (logicalExpression.charAt(i) == '&') {
                                        operations[ptr++] = "&&";
                                        i++;
                                        continue;
                                    }
                                    if (logicalExpression.charAt(i) == '|') {
                                        operations[ptr++] = "||";
                                        i++;
                                        continue;
                                    }
                                }
                                operations[ptr] = "&&";

                                for (int i = 0; i < conditions.size(); i++) {
                                    //要改变conditionInSoot，替换到logicalExpression中
                                    String conditionInSoot = conditions.get(i);
                                    String conditionInFact = logicalExpressionSplit[i];
                                    String operation = operations[i];

                                    if (operation.equals("&&")) {//||的情况下不变 &&会变成反的
                                        String operationInCondition = getOperationInCondition(conditionInSoot);
                                        String reverse = reverseOperation(operationInCondition);
                                        conditionInSoot = conditionInSoot.replace(operationInCondition, reverse);
                                    }

                                    logicalExpression = logicalExpression.replaceFirst(conditionInFact, conditionInSoot);
                                }


                                logicalExpressionSplit = logicalExpression.split("&&|\\|\\|");
                                for (int i = 0; i < logicalExpressionSplit.length; i++) {//去掉一半括号的情况
                                    String s = logicalExpressionSplit[i].trim();
                                    if (s.contains("(") && !s.contains(")")) {
                                        s = s.substring(1);
                                    }
                                    if (s.contains(")") && !s.contains("(")) {
                                        s = s.substring(0, s.length() - 1);
                                    }
                                    logicalExpressionSplit[i] = s;
                                }

                                StringBuilder preDecisions = new StringBuilder("( ");
                                for (String cond : currentStepConditions) {
                                    preDecisions.append(cond).append(" && ");
                                }
                                if (preDecisions.toString().contains("&&")) {
                                    preDecisions = new StringBuilder(preDecisions.substring(0, preDecisions.lastIndexOf("&&")));
                                }
                                preDecisions.append(")");
                                System.out.println("preDecisions: " + preDecisions);

                                // 利用z3求解
                                // preDecisions && (trueDecision !|| falseDecision) && curCondition
                                // preDecisions && (trueDecision !|| falseDecision) && !curCondition
                                for (String s : logicalExpressionSplit) {
                                    String curCondition = s;
                                    System.out.println("curCondition: " + curCondition);
                                    String trueDecision = "( " + logicalExpression.replace(curCondition, "1 == 1") + " )";
                                    System.out.println("trueDecision: " + trueDecision);
                                    String falseDecision = "( " + logicalExpression.replace(curCondition, "1 == 0") + " )";
                                    System.out.println("falseDecision: " + falseDecision);

                                    String z3Expression;
                                    if (stepConditions.size() == 0) {
                                        z3Expression = "( " + trueDecision + " !|| " + falseDecision + " )" + " && " + curCondition;
                                    } else {
                                        z3Expression = preDecisions + " && ( " + trueDecision + " !|| " + falseDecision + " )" + " && " + curCondition;
                                    }
                                    System.out.println("z3Expression: " + z3Expression);
                                    String res1 = Z3Solver.solve(z3Expression);
                                    System.out.println("res1: " + res1);

                                    String operationInCondition = getOperationInCondition(curCondition);
                                    curCondition = curCondition.replace(operationInCondition, reverseOperation(operationInCondition));

                                    if (stepConditions.size() == 0) {
                                        z3Expression = "( " + trueDecision + " !|| " + falseDecision + " )" + " && " + curCondition;
                                    } else {
                                        z3Expression = preDecisions + " && ( " + trueDecision + " !|| " + falseDecision + " )" + " && " + curCondition;
                                    }
                                    String res2 = Z3Solver.solve(z3Expression);
                                    System.out.println("res2: " + res2);

                                    if (!res1.isEmpty() && !res2.isEmpty()) {
                                        coveredCnt++;
                                        sootTestSet.add(res1);
                                        sootTestSet.add(res2);
                                    }
                                }

                            }
                            String ifstmt = ((JIfStmt) unit).getCondition().toString();
                            String cond = "( " + ifstmt + " )";

                            //替换条件里的局部变量
                            for (Local localVar : localVars) {
                                if (cond.contains(localVar.toString()) && assignList.get(localVar.toString()) != null) {
                                    cond = cond.replace(localVar.toString(), assignList.get(localVar.toString()));
                                    //stopFlag= false;
                                }
                            }
                            stepConditions.add(cond);
                        }
                    }
                }
            }
        }
        //eliminate redundancy
        sootTestSet = (ArrayList<String>) sootTestSet.stream().distinct().collect(Collectors.toList());
//        System.out.println(sootTestSet);
//        return sootTestSet;
        ArrayList<String> testSet = new ArrayList<>();
        ArrayList<String> expectedResultSet = new ArrayList<>();
        for (String res : sootTestSet) {
            if (res.equals("")) {
                testSet.add("");
                expectedResultSet.add("无解");
            } else {
//                List<Object> testObjects = getTestObjectsFromResString(res);
//                testSet.add(getTestStringFromTestObjects(testObjects));
//                Object expectedResult = getExpectedResult(testObjects);
//                expectedResultSet.add(expectedResult.toString());
            }
        }
        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            for (int i = 0; i < testSet.size(); i++) {
                System.out.println("( " + (i + 1) + " ) " + testSet.get(i));
                System.out.println("      ExpectedResult：" + expectedResultSet.get(i));
            }
        }

        coverageRate = coveredCnt / (double) conditionCnt;
        System.out.println("MCDC覆盖率：" + coverageRate);
        return testSet;
    }
    @Override
    void init() {
        PrimePathCal primePathCal = new PrimePathCal(body);
        initSet = new HashSet<>(primePathCal.generatePrimePathUnit()); //初始集合是所有的基路径
        solvableSet = new HashSet<>();
        unsolvableSet = new HashSet<>(initSet);
        testData = new HashSet<>();

    }

    @Override
    Set<List<Unit>> calAllFullCandidate(Object o) {
        List<Unit> primePath = (List)o;
        Unit headOfPrimePath = primePath.get(0);
        Unit tailOfPrimePath = primePath.get(primePath.size()-1);
        Set<List<Unit>> backwardPaths = new HashSet<>();
        Set<List<Unit>> forwardPaths = new HashSet<>();
        for (Unit head : heads) {
            PathUtil pathUtil = new PathUtil();
            pathUtil.findPath(ug, headOfPrimePath, head, new ArrayList<>(), true, new HashMap<>());
            backwardPaths.addAll(pathUtil.getSearchPathResult());
        }
        for (Unit tail : tails) {
            PathUtil pathUtil = new PathUtil();
            pathUtil.findPath(ug, tailOfPrimePath, tail, new ArrayList<>(), false, new HashMap<>());
            forwardPaths.addAll(pathUtil.getSearchPathResult());
        }
        Set<List<Unit>> result = new HashSet<>();
        if (backwardPaths.isEmpty() || forwardPaths.isEmpty()) {
            return result; //没有完整路径，返回空
        }
        for (List<Unit> backwardPath : backwardPaths) {
            Collections.reverse(backwardPath); //这个地方注意，每次得复制一下
            for (List<Unit> forwardPath : forwardPaths) {
                List<Unit> backwardPathCopy = new ArrayList<>(backwardPath);
                List<Unit> forwardPathCopy = new ArrayList<>(forwardPath); //注意不能直接操作
                backwardPathCopy.remove(backwardPathCopy.size()-1); //首尾重合了
                backwardPathCopy.addAll(primePath);
                forwardPathCopy.remove(0);
                backwardPathCopy.addAll(forwardPathCopy);//拼接到一起
                result.add(backwardPathCopy);
            }
        }
        return result;
    }

    @Override
    void checkCov(List<Unit> fullPath) {
        for (Object o : unsolvableSet) {
            List<Unit> primePath = (List<Unit>) o;
            if (CommonUtil.leftIsSubList(primePath, fullPath)) {
                solvableSet.add(primePath);
            }
        }
        unsolvableSet.removeAll(solvableSet);
    }

    private void getDecisions() {
        UnitPatchingChain units = body.getUnits();
        for (Unit unit : units) {
            if (unit instanceof JIfStmt) {
                int lineNumber = unit.getJavaSourceStartLineNumber();
                if (!decisions.containsKey(lineNumber)) {
                    ArrayList<Unit> temp = new ArrayList<>();
                    decisions.put(lineNumber, temp);
                }
                decisions.get(lineNumber).add(unit);
            }
        }
    }

    //backTrack
    private void getAllPathsHead2Tail
    (List<List<Unit>> resultPaths, List<Unit> path, HashMap<Unit, Integer> visited, Unit targetTail) {
        Unit end = path.get(path.size() - 1);
        //target found
        if (end.equals(targetTail)) {
            resultPaths.add(new ArrayList<>(path));
            return;
        }
        List<Unit> succs = ug.getSuccsOf(end);

        //path ends with failure
        if (succs == null || succs.size() == 0) {
            return;
        }
        for (Unit succ : succs) {
            if (!visited.containsKey(succ) || visited.get(succ) <= LOOP_MAX_CNT_IN_EXTEND) {
                path.add(succ);
                int cnt = 0;
                if (visited.containsKey(succ))
                    cnt = visited.get(succ);
                cnt++;
                visited.put(succ, cnt);
                getAllPathsHead2Tail(resultPaths, path, visited, targetTail);

                path.remove(path.size() - 1);
                visited.put(succ, cnt - 1);
            } else {
                if (end instanceof JIfStmt) continue; //jump out of the loop
                else {
                    path.add(succ);
                    int cnt = 0;
                    if (visited.containsKey(succ))
                        cnt = visited.get(succ);
                    cnt++;
                    visited.put(succ, cnt);
                    getAllPathsHead2Tail(resultPaths, path, visited, targetTail);

                    path.remove(path.size() - 1);
                    visited.put(succ, cnt - 1);
                }
            }
        }
    }

    public void getSplitContent(String javaPath) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(Files.newInputStream(Paths.get(javaPath))));
        String data;
        StringBuilder builder = new StringBuilder();
        while ((data = br.readLine()) != null) {
            builder.append(data).append("\n");
        }
        String content = builder.toString();
        this.splitContent = content.split("\n");
    }

    public String getLogicalExpressionByLine(int number) {
        String line = splitContent[number - 1].trim();
        if (line.contains("if") || line.contains("while")) {
            return line.substring(line.indexOf('(') + 1, line.lastIndexOf(')')).trim();
        }
        else if (line.contains("for")) {
            return line.substring(line.indexOf(';') + 1, line.lastIndexOf(';')).trim();
        }
        return "";
    }

    public static String reverseOperation(String operation) {
        if (operation.equals(">=")) {
            return "<";
        }
        if (operation.equals(">")) {
            return "<=";
        }
        if (operation.equals("<=")) {
            return ">";
        }
        if (operation.equals("<")) {
            return ">=";
        }
        if (operation.equals("==")) {
            return "!=";
        }
        if (operation.equals("!=")) {
            return "==";
        }
        return "unknown";
    }

    public static String getOperationInCondition(String condition) {
        if (condition.contains(">=")) {
            return ">=";
        }
        if (condition.contains("<=")) {
            return "<=";
        }
        if (condition.contains("==")) {
            return "==";
        }
        if (condition.contains(">")) {
            return ">";
        }
        if (condition.contains("<")) {
            return "<";
        }
        if (condition.contains("!=")) {
            return "!=";
        }
        return "unknown";
    }

}
